from __future__ import division

import logging, math

import numpy, scipy.interpolate, scipy.optimize

import common

logger = logging.getLogger("wlc")
logger.addHandler(logging.StreamHandler())
logger.setLevel(logging.INFO)

# "Contour Length and Refolding Rate of a Small Protein Controlled by Engineered Disulfide Bonds" uses instead 3 Angstrom.
default_p = 0.35e-9

default_contour_per_residue = 0.4e-9

def f_derived(kb_t_over_p, k_lc, v_over_lc, t):
    x_over_lc = find_x_over_lc_for_tug_of_war(kb_t_over_p, k_lc, v_over_lc * t)
    y = wlc(kb_t_over_p, x_over_lc)
    return 1 / f_inverse_derived(kb_t_over_p, k_lc, v_over_lc, y)

# Finds the time!
def f_inverse(kb_t_over_p, k_lc, v_over_lc, y):
    return core_f_inverse(kb_t_over_p, k_lc, v_over_lc, y)[0]

def core_f_inverse(kb_t_over_p, k_lc, v_over_lc, y):
    """
       TODO: Would it be better/faster to use a hand-coded cubic equation analytic solver?
       Return: numpy.bestfloat64
    """
    assert y >= 0
    # assert isinstance(kb_t_over_p, float)
    # assert isinstance(k_lc, float)
    # assert isinstance(v_over_lc, float)
    # assert isinstance(y, float)

    kb_t_over_p_over_k_lc = kb_t_over_p / k_lc
    twelve_kb_t_over_p_over_k_lc_plus_four_times_y = (
        12 * kb_t_over_p_over_k_lc + 4) * y
    six_kb_t_over_p_over_k_lc_plus_four_times_y = (
        6 * kb_t_over_p_over_k_lc + 4) * y
    y_square = y ** 2

    c0 = -4 * kb_t_over_p
    c1 = twelve_kb_t_over_p_over_k_lc_plus_four_times_y + 9 * kb_t_over_p
    c2 = -((twelve_kb_t_over_p_over_k_lc_plus_four_times_y + 4 * y) / k_lc * y + (
            twelve_kb_t_over_p_over_k_lc_plus_four_times_y + six_kb_t_over_p_over_k_lc_plus_four_times_y) + 6 * kb_t_over_p)
    # Note that numpy.power(k_lc, 2) works for k_lc=inf, while inf ** 2 fails.
    c3a = (4 * kb_t_over_p_over_k_lc + 4) / numpy.power(k_lc, 2) * y_square * y
    c3b = (9 * kb_t_over_p_over_k_lc + 8) / k_lc * y_square
    c3 = c3a + c3b + six_kb_t_over_p_over_k_lc_plus_four_times_y
    c = c0, c1, c2, c3,

    # Numpy.roots is the fastest of all.
    roots = numpy.roots(c)
    # roots = cubic.real_roots(*c)

    non_negative_roots = list(frozenset([r.real for r in roots if r.imag == 0 and r.real >= 0]))
    _ts = tuple(r / v_over_lc for r in non_negative_roots)
    # Select the unique root giving 0 <= x(t) / Lc < 1
    ts = tuple(t for t in _ts if 0 <= v_over_lc * t - y / k_lc < 1)
    if len(ts) != 1:
        print(list(t - y / k_lc for t in non_negative_roots), roots)
        print("Coefficients: %s" % (c, ))
        print("Expected: %s" % numpy.roots(c))
        print("Actual: %s" % (roots, ))
        raise Exception("Could not find exactly one good root for y=%s coefficients:%s. Instead I found roots:%s and fractions %r." % (y, c, roots, [v_over_lc * t - y / k_lc for t in _ts]))
    r = ts[0]
    return r, c

def wlc_inverse_derived_times_v_over_l_minus_1_over_kl(kb_t_over_p, y):
    return 2 / kb_t_over_p * (1 + y / kb_t_over_p) / (
        3 + 5 * y / kb_t_over_p + 8 * (y / kb_t_over_p) ** 2.5)

def f_inverse_derived_approximation(kb_t_over_p, k_lc, v_over_lc, y):
    """
    Dudko OK, Hummer G, Szabo A.
    "Theory, analysis, and interpretation of single-molecule force spectroscopy experiments",
    Proc Natl Acad Sci U S A. 2008 Oct 14;105(41):15755-60
    """
    return (
        1 / k_lc + wlc_inverse_derived_times_v_over_l_minus_1_over_kl(kb_t_over_p, y)
        ) / v_over_lc

def f_inverse_derived(kb_t_over_p, k_lc, v_over_lc, y):
    t, coeffs = core_f_inverse(kb_t_over_p, k_lc, v_over_lc, y)
    z = v_over_lc * t
    kb_t_over_p_over_k_lc = kb_t_over_p / k_lc
    twelve_kb_t_over_p_over_k_lc_plus_4 = 12 * kb_t_over_p_over_k_lc + 4
    six_kb_t_over_p_over_k_lc_plus_4 = 6 * kb_t_over_p_over_k_lc + 4
    eighteen_kb_t_over_p_over_k_lc_plus_8 = twelve_kb_t_over_p_over_k_lc_plus_4 + six_kb_t_over_p_over_k_lc_plus_4
    db_over_df = twelve_kb_t_over_p_over_k_lc_plus_4
    dc_over_df = -((2 * twelve_kb_t_over_p_over_k_lc_plus_4 + 8) / k_lc * y + eighteen_kb_t_over_p_over_k_lc_plus_8)
    dd_over_df = (twelve_kb_t_over_p_over_k_lc_plus_4 + 8) / k_lc ** 2 * y ** 2 + (
        eighteen_kb_t_over_p_over_k_lc_plus_8 + 8) / k_lc * y + six_kb_t_over_p_over_k_lc_plus_4
    z_square = z ** 2
    dz_over_df = (-db_over_df * z_square - dc_over_df * z - dd_over_df) / (
        3 * coeffs[0] * z_square + 2 * coeffs[1] * z + coeffs[2])
    result = dz_over_df / v_over_lc
    return result

def wlc(kb_t_over_p, x_over_lc):
    if numpy.isscalar(x_over_lc):
        if x_over_lc == 1:
            return numpy.inf
    # TODO: Avoid zero division for numpy arrays as well.
    return kb_t_over_p * (1.0 / 4.0 / (1.0 - x_over_lc) ** 2 + x_over_lc - 1.0 / 4.0)

def wlc_indefinite_integral(kb_t_over_p, x_over_lc):
    common.check_le(0, x_over_lc)
    common.check_le(x_over_lc, 1)
    return kb_t_over_p * (1 / 4 / (1 - x_over_lc) + x_over_lc * (
            x_over_lc / 2 - 1 / 4))

def wlc_definite_integral(kb_t_over_p, a, b):
    return wlc_indefinite_integral(kb_t_over_p, b) - wlc_indefinite_integral(
        kb_t_over_p, a)

def compute_x_over_lc(kb_t_over_p, y):
    return f_inverse(kb_t_over_p, numpy.inf, 1, y)

def find_x_for_tug_of_war(kb_t_over_p, lc, k, z):
    return find_x_over_lc_for_tug_of_war(kb_t_over_p, k * lc, z / lc) * lc

def find_x_over_lc_for_tug_of_war(kb_t_over_p, k_lc, z_over_lc):
    # Around 90 ms.
    root, r = scipy.optimize.brentq(
        lambda x_over_lc: wlc(kb_t_over_p, x_over_lc) - k_lc * (
            z_over_lc - x_over_lc), 0, 1, full_output=True, xtol=1e-17)
    common.check_True(r.converged)
    return root

def radial_probability_density(k, r):
    """
    k = l_p / l
    r = distance / l
    "The radial distribution function of worm-like chains",
    EUR PHYS J E , vol. 32, no. 2, 2010
    """
    common.check_True((numpy.asarray(r) >= 0).all())
    common.check_True((numpy.asarray(r) <= 1).all())
    one_minus_r2 = 1 - r ** 2
    if not one_minus_r2:
        return 0
    a, b = 14.054, 0.473
    c = 1 - (1 + (0.38 * k ** (-0.95)) ** (-5)) ** (-1 / 5)
    d = 1 if k < 1 / 8 else 1 - 1 / (
        0.177 / (k - 0.111) + 6.40 * (k - 0.111) ** 0.783)
    x = -d * k * a * (1 + b) * r / (1 - b ** 2 * r ** 2)
    # Note that a in j_syd (section 3.2) is the same as a in q_i.
    j_syd = (
        112.04 * k ** 2 * numpy.exp(0.246 / k - a * k) if k > 1 / 8
        else (3 / (4 * numpy.pi * k)) ** (3 / 2)  * (1 - 5 * k / 4))
    C = ((-3 / 4, 23 / 64, -7 / 64), (-1 / 2, 17 / 16, -9 / 16))
    logger.debug("r=%s k=%s x=%s a=%s b=%s d=%s", r, k, x, a, b, d)
    q_i = j_syd * ((1 - c * r ** 2) / one_minus_r2) ** (5 / 2) * numpy.exp(sum(
            C[i + 1][j - 1] * k ** i * r ** (2 * j) for i in range(-1, 1)
            for j in range(1, 4)) / one_minus_r2 + b * r * x) * numpy.i0(-x)
    return 4 * numpy.pi * r**2 * q_i

def log_radial_probability_density(k, r):
    common.check_True(numpy.isscalar(k))
    if numpy.isscalar(r):
        if r == 0 or r == 1:
            return -numpy.inf
        else:
            non_zero_one_r = numpy.asarray([r], dtype=numpy.longfloat)
    else:
        r_array = numpy.asarray(r, dtype=numpy.longfloat)
        non_zero_one_r = r_array[(r_array != 0) * (r_array != 1)]
    if len(non_zero_one_r) == 0:
        rpd = numpy.asarray([], dtype=numpy.longfloat)
    else:
        rpd = log_radial_probability_density_non_zero_one_r_array(k, non_zero_one_r)
    if numpy.isscalar(r):
        common.check_equal(len(rpd), 1)
        result = rpd[0]
    else:
        result = numpy.asarray([-numpy.inf] * len(r), dtype=numpy.longfloat)
        result[(r_array != 0) * (r_array != 1)] = rpd
    return result
        
def log_radial_probability_density_non_zero_one_r_array(k, non_zero_one_r):
    common.check_True((non_zero_one_r > 0).all(), non_zero_one_r)
    common.check_True((non_zero_one_r < 1).all(), non_zero_one_r)
    one_minus_r2 = 1 - non_zero_one_r ** 2
    a, b = 14.054, 0.473
    c = 1 - (1 + (0.38 * k ** (-0.95)) ** (-5)) ** (-1 / 5)
    d = 1 if k < 1 / 8 else 1 - 1 / (
        0.177 / (k - 0.111) + 6.40 * (k - 0.111) ** 0.783)
    x = -d * k * a * (1 + b) * non_zero_one_r / (1 - b ** 2 * non_zero_one_r ** 2)
    log_j_syd = (
        numpy.log(112.04) + numpy.log(k) * 2 + 0.246 / k - a * k if k > 1 / 8
        else numpy.log(3 / (4 * numpy.pi * k)) * (3 / 2) + numpy.log(1 - 5 * k / 4))
    C = ((-3 / 4, 23 / 64, -7 / 64), (-1 / 2, 17 / 16, -9 / 16))
    log_q_i = log_j_syd + numpy.log((1 - c * non_zero_one_r ** 2) / one_minus_r2) * (5 / 2) + sum(
            C[i + 1][j - 1] * k ** i * non_zero_one_r ** (2 * j) for i in range(-1, 1)
            for j in range(1, 4)) / one_minus_r2 + b * non_zero_one_r * x + numpy.log(
        numpy.i0(-x))
    return numpy.log(4 * numpy.pi) + 2 * numpy.log(non_zero_one_r) + log_q_i
